// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <tests/expression/expression_runner.h>
#include <melon/init/init.h>
#include <turbo/flags/flag.h>
#include <gtest/gtest.h>
#include <pollux/common/base/fs.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/connectors/hive/hive_connector.h>
#include <pollux/dwio/dwrf/register_dwrf_writer.h>
#include <pollux/testing/exec/fuzzer_util.h>
#include <pollux/testing/flags.h>
#include <pollux/testing/exec/presto_query_runner.h>
#include <pollux/plan/reference_query_runner.h>
#include <pollux/testing/expression/expression_verifier.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/functions/sparksql/registration/register.h>
#include <pollux/vector/vector_saver.h>

using namespace kumo::pollux;
using namespace kumo::pollux::plan;
using kumo::pollux::exec::test::PrestoQueryRunner;

TURBO_FLAG(std::string,
           input_paths,
           "",
           "Comma separated list of paths for vectors to be restored from disk. This "
           "will enable single run of the fuzzer with the on-disk persisted repro "
           "information. This has to be set with sql_path and optionally "
           "result_path.");

TURBO_FLAG(std::string,
           input_selectivity_vector_paths,
           "",
           "Comma separated list of paths for selectivity vectors to be restored "
           "from disk. The list needs to match 1-to-1 with the files specified in "
           "input_paths. If not specified, all rows will be selected.");

TURBO_FLAG(std::string,
           sql_path,
           "",
           "Path for expression SQL to be restored from disk. This will enable "
           "single run of the fuzzer with the on-disk persisted repro information. "
           "This has to be set with input_path and optionally result_path.");

TURBO_FLAG(std::string,
           complex_constant_path,
           "",
           "Path for complex constants that aren't expressible in SQL.");

TURBO_FLAG(std::string,
           sql,
           "add(2, 3), ceil(1.3)",
           "Comma separated SQL expressions to evaluate. This flag and --sql_path "
           "flag are mutually exclusive. If both are specified, --sql is used and "
           "--sql_path is ignored.");

TURBO_FLAG(std::string,
           registry,
           "presto",
           "Funciton registry to use for expression evaluation. Currently supported values are "
           "presto and spark. Default is presto.");

TURBO_FLAG(std::string,
           result_path,
           "",
           "Path for result vector to restore from disk. This is optional for "
           "on-disk reproduction. Don't set if the initial repro result vector is "
           "nullptr");

TURBO_FLAG(std::string,
           mode,
           "common",
           "Mode for expression runner: \n"
           "verify: evaluate the expression and compare results between common and "
           "simplified paths.\n"
           "common: evaluate the expression using common path and print out results.\n"
           "simplified: evaluate the expression using simplified path and print out "
           "results.\n"
           "query: evaluate SQL query specified in --sql or --sql_path and print out "
           "results. If --input_paths is specified, the query may reference it as "
           "table 't'. Note: --input_selectivity_vector_paths is ignored in this mode");

TURBO_FLAG(std::string,
           input_row_metadata_path,
           "",
           "Path for the file stored on-disk which contains a struct containing "
           "input row metadata. This includes columns in the input row vector to "
           "be wrapped in a lazy vector and/or dictionary encoded. It may also "
           "contain a dictionary peel for columns requiring dictionary encoding.");

TURBO_FLAG(bool,
           use_seperate_memory_pool_for_input_vector,
           true,
           "If true, expression evaluator and input vectors use different memory pools."
           " This helps trigger code-paths that can depend on vectors having different"
           " pools. For eg, when copying a flat string vector copies of the strings"
           " stored in the string buffers need to be created. If however, the pools"
           " were the same between the vectors then the buffers can simply be shared"
           " between them instead.");

TURBO_FLAG(std::string,
           reference_db_url,
           "",
           "ReferenceDB URI along with port. If set, we use the reference DB as the "
           "source of truth. Otherwise, use Pollux simplified eval path. Example: "
           "--reference_db_url=http://127.0.0.1:8080");

TURBO_FLAG(uint32_t,
           req_timeout_ms,
           10000,
           "Timeout in milliseconds for HTTP requests made to reference DB, "
           "such as Presto. Example: --req_timeout_ms=2000");

static bool validateMode(const char *flagName, const std::string &value) {
    static const std::unordered_set<std::string> kModes = {
        "common", "simplified", "verify", "query"
    };
    if (kModes.count(value) != 1) {
        std::cerr << "Invalid value for --" << flagName << ": " << value << ". ";
        std::cerr << "Valid values are: " << melon::join(", ", kModes) << "."
                << std::endl;
        return false;
    }

    return true;
}

static bool validateRegistry(const char *flagName, const std::string &value) {
    static const std::unordered_set<std::string> kRegistries = {
        "presto", "spark"
    };
    if (kRegistries.count(value) != 1) {
        std::cerr << "Invalid value for --" << flagName << ": " << value << ". ";
        std::cerr << "Valid values are: " << melon::join(", ", kRegistries) << "."
                << std::endl;
        return false;
    }
    if (value == "spark") {
        functions::sparksql::registerFunctions("");
    } else if (value == "presto") {
        functions::prestosql::registerAllScalarFunctions();
    }

    return true;
}

//DEFINE_validator(mode, &validateMode);
//DEFINE_validator(registry, &validateRegistry);

TURBO_FLAG(int32_t,
           num_rows,
           10,
           "Maximum number of rows to process. Zero means 'all rows'. Applies to "
           "'common' and 'simplified' modes only. Ignored for 'verify' mode.");

TURBO_FLAG(std::string,
           store_result_path,
           "",
           "Directory path for storing the results of evaluating SQL expression or "
           "query in common, simplified or query modes.");

TURBO_FLAG(std::string,
           fuzzer_repro_path,
           "",
           "Directory path where all input files generated by ExpressionVerifier are "
           "expected to reside. For more details on which files and their names are "
           "expected, please checkout the ExpressionVerifier class. Any file paths "
           "already specified via a startup flag will take precedence.");


static std::string checkAndReturnFilePath(
    const std::string_view &fileName,
    const std::string &flagName) {
    auto path = fmt::format("{}/{}", turbo::get_flag(FLAGS_fuzzer_repro_path), fileName);
    if (fs::exists(path)) {
        KLOG(INFO) << "Using " << flagName << " = " << path;
        return path;
    } else {
        KLOG(INFO) << "File for " << flagName << " not found.";
    }
    return "";
}

static std::string getFilesWithPrefix(
    const char *dirPath,
    const std::string_view &prefix,
    const std::string &flagName) {
    std::vector<std::string> filesPaths;
    std::stringstream ss;
    int numFilesFound = 0;
    if (!std::filesystem::exists(dirPath)) {
        KLOG(ERROR) << "Directory does not exist: " << dirPath << std::endl;
        return "";
    }
    for (const auto &entry: std::filesystem::directory_iterator(dirPath)) {
        if (entry.is_regular_file()) {
            std::string filename = entry.path().filename();
            if (filename.find(prefix) == 0) {
                if (++numFilesFound > 1) {
                    ss << ",";
                }
                ss << entry.path().string();
            }
        }
    }
    KLOG(INFO) << "Using " << flagName << " = " << ss.str();
    return ss.str();
}

static void checkDirForExpectedFiles() {
    KLOG(INFO) << "Searching input directory for expected files at "
            << turbo::get_flag(FLAGS_fuzzer_repro_path);

    turbo::set_flag(&FLAGS_input_paths, turbo::get_flag(FLAGS_input_paths).empty()
                                            ? getFilesWithPrefix(
                                                turbo::get_flag(FLAGS_fuzzer_repro_path).c_str(),
                                                test::ExpressionVerifier::kInputVectorFileNamePrefix,
                                                "input_paths")
                                            : turbo::get_flag(FLAGS_input_paths));
    turbo::set_flag(&FLAGS_input_selectivity_vector_paths,
                    turbo::get_flag(FLAGS_input_selectivity_vector_paths).empty()
                        ? getFilesWithPrefix(
                            turbo::get_flag(FLAGS_fuzzer_repro_path).c_str(),
                            test::ExpressionVerifier::kInputSelectivityVectorFileNamePrefix,
                            "input_selectivity_vector_paths")
                        : turbo::get_flag(FLAGS_input_selectivity_vector_paths));
    turbo::set_flag(&FLAGS_result_path, turbo::get_flag(FLAGS_result_path).empty()
                                            ? checkAndReturnFilePath(
                                                test::ExpressionVerifier::kResultVectorFileName, "result_path")
                                            : turbo::get_flag(FLAGS_result_path));
    turbo::set_flag(&FLAGS_sql_path, turbo::get_flag(FLAGS_sql_path).empty()
                                         ? checkAndReturnFilePath(
                                             test::ExpressionVerifier::kExpressionSqlFileName, "sql_path")
                                         : turbo::get_flag(FLAGS_sql_path));
    turbo::set_flag(&FLAGS_input_row_metadata_path, turbo::get_flag(FLAGS_input_row_metadata_path).empty()
                                                        ? checkAndReturnFilePath(
                                                            test::ExpressionVerifier::kInputRowMetadataFileName,
                                                            "input_row_metadata_path")
                                                        : turbo::get_flag(FLAGS_input_row_metadata_path));
    turbo::set_flag(&FLAGS_complex_constant_path, turbo::get_flag(FLAGS_complex_constant_path).empty()
                                                      ? checkAndReturnFilePath(
                                                          test::ExpressionVerifier::kComplexConstantsFileName,
                                                          "complex_constant_path")
                                                      : turbo::get_flag(FLAGS_complex_constant_path));
}

int main(int argc, char **argv) {
    ::testing::InitGoogleTest(&argc, argv);
    // Calls common init functions in the necessary order, initializing
    // singletons, installing proper signal handlers for better debugging
    // experience
    melon::Init init(&argc, &argv);

    if (!turbo::get_flag(FLAGS_fuzzer_repro_path).empty()) {
        checkDirForExpectedFiles();
    }

    if (turbo::get_flag(FLAGS_sql).empty() && turbo::get_flag(FLAGS_sql_path).empty()) {
        std::cerr << "One of --sql or --sql_path flags must be set." << std::endl;
        exit(1);
    }

    auto sql = turbo::get_flag(FLAGS_sql);
    if (sql.empty()) {
        sql = restoreStringFromFile(turbo::get_flag(FLAGS_sql_path).c_str());
        POLLUX_CHECK(!sql.empty());
    }

    memory::initializeMemoryManager({});

    filesystems::registerLocalFileSystem();
    connector::registerConnectorFactory(
        std::make_shared<connector::hive::HiveConnectorFactory>());
    exec::test::registerHiveConnector({});
    dwrf::registerDwrfWriterFactory();

    std::shared_ptr<kumo::pollux::memory::MemoryPool> rootPool{
        kumo::pollux::memory::memoryManager()->addRootPool()
    };
    std::shared_ptr<ReferenceQueryRunner> referenceQueryRunner{nullptr};
    if (turbo::get_flag(FLAGS_registry) == "presto" && !turbo::get_flag(FLAGS_reference_db_url).empty()) {
        referenceQueryRunner = std::make_shared<PrestoQueryRunner>(
            rootPool.get(),
            turbo::get_flag(FLAGS_reference_db_url),
            "expression_runner_test",
            static_cast<std::chrono::milliseconds>(turbo::get_flag(FLAGS_req_timeout_ms)));
        KLOG(INFO) << "Using Presto as the reference DB.";
    }

    test::ExpressionRunner::run(
        turbo::get_flag(FLAGS_input_paths),
        turbo::get_flag(FLAGS_input_selectivity_vector_paths),
        sql,
        turbo::get_flag(FLAGS_complex_constant_path),
        turbo::get_flag(FLAGS_result_path),
        turbo::get_flag(FLAGS_mode),
        turbo::get_flag(FLAGS_num_rows),
        turbo::get_flag(FLAGS_store_result_path),
        turbo::get_flag(FLAGS_input_row_metadata_path),
        referenceQueryRunner,
        turbo::get_flag(FLAGS_find_minimal_subexpression),
        turbo::get_flag(FLAGS_use_seperate_memory_pool_for_input_vector));
}
